/**
 * @file heap_parse.cc
 * @brief This file contains functions for parsing and analyzing heap memory allocations.
 *
 * The functions in this file parse and analyze heap memory allocations using information from
 * the AllocRecord, SmapsData, and ThreadRecord structures. It provides functions to check if
 * an allocation record belongs to the glibc mmap region, update the allocation size for a thread,
 * parse an AllocRecord entry and perform necessary operations based on its in_use value, print
 * stack information for each stack frame, print thread allocation information, print the heap result,
 * print the mmap result, print the Smaps analysis result, and clean the linked lists.
 *
 * The main entry point for the analysis is the `PrintHeapResult` function, which prints the total
 * memory usage and allocation information for the heap. It also sorts the heap list by length in
 * descending order and prints each allocation record, including the address, size, percentage of
 * total heap size, and associated thread name. Additionally, it prints the stack information for
 * each allocation record.
 */

#include <libgen.h>
#include <zlib.h>

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <memory>
#include <sstream>
#include <vector>

#include "alloc_hook.h"

struct SmapsData {
  void* address_start;      // start address of the memory
  void* address_end;        // end address of the memory
  char mode[8];             // mode of the memory
  char name[64];            // name of the memory
  uint32_t size;            // size of the memory
  uint32_t rss;             // resident set size
  char lib_path[512];       // path of the library
  uint8_t is_lib_found;     // flag to indicate if the library is found
  uint8_t is_thread_stack;  // flag to indicate if the memory is a thread stack
  uint8_t is_thread_heap;   // flag to indicate if the memory is a thread heap
  SmapsData* next;          // next memory record
  SmapsData()
      : address_start(nullptr),
        address_end(nullptr),
        size(0),
        rss(0),
        is_lib_found(0),
        is_thread_stack(0),
        is_thread_heap(0),
        next(nullptr) {
    memset(mode, 0, sizeof(mode));
    memset(name, 0, sizeof(name));
    memset(lib_path, 0, sizeof(lib_path));
  }
};

static AllocRecord* heap_list = nullptr;
static AllocRecord* mmap_list = nullptr;
static SmapsData* smaps_list = nullptr;
static AllocDataHeader header{};
static SmapsData* elf_text = nullptr;
static ThreadRecord* thread_info = nullptr;
static char* addr2line = nullptr;
static char* elf_file = nullptr;
static char* result_file = nullptr;
static char* lib_dir = nullptr;

constexpr uint32_t kPrintStackLen = 4;

static const char* GetThreadNameByPid(const pid_t pid) {
  for (ThreadRecord* p = thread_info; p != nullptr; p = p->next) {
    if (p->pid == pid) {
      return p->name;
    }
  }
  return nullptr;
}

/**
 * Checks if the given allocation record belongs to the glibc mmap region.
 * If the allocation record is found within the glibc mmap region, it updates the corresponding SmapsData structure.
 * If the SmapsData structure does not have a name, it assigns a name to it based on the thread name associated with the
 * allocation record's process ID.
 *
 * @param entry The allocation record to be checked.
 */
static void CheckIfGlibcMmap(const AllocRecord* entry) {
  for (SmapsData* p = smaps_list; p != nullptr; p = p->next) {
    if (!p->is_thread_stack && (entry->p >= p->address_start && entry->p <= p->address_end)) {
      p->is_thread_heap = 1;
      if (strlen(p->name) == 0) {
        snprintf(p->name, sizeof(p->name), "(heap)(%s)", GetThreadNameByPid(entry->pid));
      }
      return;
    }
  }
}

/**
 * @brief Updates the allocation size for the thread associated with the given entry.
 *
 * This function iterates through the linked list of thread records and finds the thread
 * with the same process ID as the given entry. It then updates the allocation size of
 * that thread by adding the length of the entry.
 *
 * @param entry The allocation record for which the thread allocation size needs to be updated.
 */
static void CountThreadAllocSize(const AllocRecord* entry) {
  for (ThreadRecord* p = thread_info; p != nullptr; p = p->next) {
    if (p->pid == entry->pid) {
      p->alloc_size += entry->len;
      break;
    }
  }
}

/**
 * Parses an AllocRecord entry and performs necessary operations based on its in_use value.
 *
 * @param entry The AllocRecord entry to be parsed.
 */
static void AllocRecordParse(const AllocRecord* entry) {
  auto node = std::make_unique<AllocRecord>();
  memcpy(node.get(), entry, sizeof(AllocRecord));

  if (entry->in_use == ALLOC_IN_USE) {
    CheckIfGlibcMmap(entry);
    CountThreadAllocSize(entry);
    for (AllocRecord* p = heap_list; p != nullptr; p = p->next) {
      if ((memcmp(p->stack, entry->stack, sizeof(entry->stack[0]) * std::min(entry->stack_len, kPrintStackLen)) == 0) &&
          p->pid == entry->pid) {
        p->len += entry->len;
        return;
      }
    }
    node->next = heap_list;
    heap_list = node.release();
  } else if (entry->in_use == MMAP_IN_USE) {
    node->next = mmap_list;
    mmap_list = node.release();
  } else {
    printf("unknown type %d\n", entry->in_use);
  }
}

/**
 * Prints stack information for each stack frame in the given stack.
 *
 * @param stack An array of void pointers representing the stack frames.
 * @param stack_len The length of the stack array.
 */
static void PrintStackInfo(void** stack, uint32_t stack_len) {
  char cmd[1024] = {0};
  char func[1024] = {0};
  char file[1024] = {0};
  FILE* fp = nullptr;

  for (int i = 1; i < static_cast<int>(stack_len); ++i) {
    if (stack[i] > elf_text->address_start && stack[i] < elf_text->address_end) {
      snprintf(cmd, sizeof(cmd), "%s -f --demangle -C -e %s %p", addr2line, elf_file, stack[i]);
    } else {
      SmapsData* lib = smaps_list;
      while (lib != nullptr) {
        if ((stack[i] > lib->address_start) && (stack[i] < lib->address_end)) {
          int offset =
              reinterpret_cast<uintptr_t>(stack[i]) - reinterpret_cast<uintptr_t>(lib->address_start);  // NOLINT
          snprintf(cmd, sizeof(cmd), "%s -f --demangle -C -e %s 0x%x", addr2line, lib->lib_path, offset);
          break;
        }
        lib = lib->next;
      }
      if (lib == nullptr) {
        printf("# %-2d %p\n", i, stack[i]);
        continue;
      } else if (lib->is_lib_found == 0) {
        printf("# %-2d %p   %s\n", i, stack[i], lib->name);
        continue;
      }
    }
    fp = popen(cmd, "r");
    if (fp != nullptr) {
      memset(func, 0, sizeof(func));
      memset(file, 0, sizeof(file));
      if (fgets(func, sizeof(func), fp) == nullptr) {
        pclose(fp);
        continue;
      }
      if (fgets(file, sizeof(file), fp) == nullptr) {
        pclose(fp);
        continue;
      }
      pclose(fp);
      printf("# %-2d %s     %s", i, func, file);
    }
  }
  printf("\n");
}

/**
 * @brief Print thread alloc info, including thread name and alloc size
 */
static void PrintThreadAllocInfo() {
  for (ThreadRecord* p = thread_info; p != nullptr; p = p->next) {
    if (p->alloc_size) {
      printf("%16s  %8u kB\n", p->name, p->alloc_size / 1024);
    }
  }
}

/**
 * @brief Prints the heap result, including total memory usage and allocation information.
 * Sorts the heap list by length in descending order and prints each allocation record.
 * The output includes the address, size, percentage of total heap size, and associated thread name.
 * Additionally, it prints the stack information for each allocation record.
 */
static void PrintHeapResult() {
  printf("heap total memory: %u kB, alloc_hook use %u kB\n", header.heap_size / 1024, header.alloc_list_size / 1024);
  PrintThreadAllocInfo();

  // Sort the heap list by length in descending order
  std::vector<AllocRecord*> heap_vector;
  for (AllocRecord* p = heap_list; p != nullptr; p = p->next) {
    heap_vector.push_back(p);
  }
  std::sort(heap_vector.begin(), heap_vector.end(),
            [](const AllocRecord* a, const AllocRecord* b) { return b->len < a->len; });

  for (AllocRecord* p1 : heap_vector) {
    if (p1->len * 100 / header.heap_size == 0) {
      printf(".");
      continue;
    }
    printf("%p  %u kB  %u.%u%%   (%s)\n", p1->p, p1->len / 1024, p1->len * 100 / header.heap_size,
           p1->len * 10000 / header.heap_size % 100, GetThreadNameByPid(p1->pid));
    PrintStackInfo(p1->stack, std::min(p1->stack_len, kPrintStackLen));
  }
}

/**
 * @brief Prints the mmap result.
 * This function prints the total size of mmap in kilobytes and then iterates over the mmap_list
 * to print the address and length of each mmap allocation. It also calculates the RSS (Resident Set Size)
 * for each mmap allocation by summing up the rss values from the smaps_list. Finally, it calls the
 * PrintStackInfo function to print the stack information for each mmap allocation.
 */
static void PrintMmapResult() {
  printf("\nmmap total size: %u kB\n", header.mmap_size / 1024);
  for (AllocRecord* p = mmap_list; p != nullptr; p = p->next) {
    printf("mmap %p len 0x%x\n", p->p, p->len);
    uint32_t rss = 0;
    for (SmapsData* p1 = smaps_list; p1 != nullptr; p1 = p1->next) {
      if ((p->p == p1->address_start) ||
          ((reinterpret_cast<uintptr_t>(p->p) + p->len) == reinterpret_cast<uintptr_t>(p1->address_end))) {
        rss += p1->rss;
      }
    }
    printf("RSS: %u kB\n", rss);
    PrintStackInfo(p->stack, p->stack_len);
  }
}

/**
 * @brief Prints the result of the Smaps analysis.
 * This function iterates through the SmapsData linked list and prints the address range,
 * mode, RSS (Resident Set Size), and name of each entry. It also calculates the total RSS
 * and the RSS of the libraries.
 */

static void PrintSmapsResult() {
  uint32_t lib_rss = 0;
  uint32_t total_rss = 0;
  for (SmapsData* p = smaps_list; p != nullptr; p = p->next) {
    if (p->rss) {
      const char* name = (strlen(p->name) == 0) ? "[anonymous]" : basename(p->name);
      printf("%12p-%12p  %-6s RSS: %8u kB  %-30s\n", p->address_start, p->address_end, p->mode, p->rss, name);
      if (strstr(name, "lib") && strstr(name, ".so")) {
        lib_rss += p->rss;
      }
    }
    total_rss += p->rss;
  }
  printf("total RSS %ukB, lib RSS %u kB, text RSS %u kB\n", total_rss, lib_rss, elf_text->rss);
}

/**
 * @brief Clean the list
 */
static void CleanList() {
  while (heap_list) {
    AllocRecord* p = heap_list;
    heap_list = heap_list->next;
    free(p);
  }
  while (mmap_list) {
    AllocRecord* p = mmap_list;
    mmap_list = mmap_list->next;
    free(p);
  }
  while (smaps_list) {
    SmapsData* p = smaps_list;
    smaps_list = smaps_list->next;
    free(p);
  }
}

/**
 * @brief Set stack info
 * traverses a linked list of memory map entries,
 * identifies those without a name and within the stack range of any thread,
 * and labels them as thread stacks.
 */
static void SetStackInfo() {
  for (SmapsData* smaps_entry = smaps_list; smaps_entry != nullptr; smaps_entry = smaps_entry->next) {
    // Check if the memory is a thread stack
    if (strlen(smaps_entry->name) == 0) {
      // Traverse the thread info list
      for (ThreadRecord* thread_entry = thread_info; thread_entry != nullptr; thread_entry = thread_entry->next) {
        if (thread_entry->kstk_esp >= smaps_entry->address_start &&
            thread_entry->kstk_esp <= smaps_entry->address_end) {
          snprintf(smaps_entry->name, sizeof(smaps_entry->name), "(stack)(%s)", thread_entry->name);
          smaps_entry->is_thread_stack = 1;
          break;
        }
      }
    }
  }
}

/**
 * @brief Strip space
 * @param str the string to be stripped
 */
static void StripSpace(char* str) {
  char* p = str;
  while (*p == ' ') {
    ++p;
  }
  memmove(str, p, strlen(p) + 1);
}

/**
 * @brief This function reads fields from a stream.
 * @param stream The file stream.
 * @param format The format string.
 * @param... The variable arguments.
 * @return true if the fields are successfully read, false otherwise.
 */
static bool FscanfReadFields(FILE* stream, const char* format, const char* const field_name, int expected_count, ...) {
  va_list args;
  va_start(args, expected_count);
  int result = vfscanf(stream, format, args);
  va_end(args);
  (void)field_name;

  // If the number of successfully read fields is not equal to the number of expected fields.
  if (result != expected_count) {  // Check the number of successfully read fields
    if (feof(stream) || ferror(stream)) {
      clearerr(stream);
      // printf("read '%s' error, break. \n", field_name);
      return false;  // File end or read error
    } else {
      // printf("failed to read '%s'. \n", field_name);
      return true;  // Failed to read as expected, but return true
    }
  }
  return true;
}

int main(int argc, char** argv) {
  if (argc < 4) {
    printf("heap_parse [addr2line] [elf_file] [heap_file] [lib_dir]\n");
    return -1;
  }
  addr2line = argv[1];
  elf_file = argv[2];
  result_file = argv[3];
  if (argc == 5) {
    lib_dir = argv[4];
  }

  // Open the result file
  auto infile = gzopen(result_file, "rb");
  if (!infile) {
    printf("can't gzopen %s\n", result_file);
    return -1;
  }

  //
  const char* const tmp_file = "/tmp/heap_data";
  FILE* fp = fopen(tmp_file, "w+");
  if (!fp) {
    gzclose(infile);
    printf("open %s failed. %m\n", tmp_file);
    return -1;
  }

  // Read the result file
  char buff[1024]{};
  int len{};
  while ((len = gzread(infile, buff, sizeof(buff))) > 0) {
    if (fwrite(buff, 1, len, fp) != len) {
      fprintf(stderr, "failed fwrite");
    }
  }
  if (len < 0) {
    printf("gzread failed.\n");
    gzclose(infile);
    fclose(fp);
    unlink(tmp_file);
    return -1;
  }
  gzclose(infile);

  fseek(fp, 0, SEEK_SET);
  if (fread(&header, 1, sizeof(header), fp) != sizeof(header)) {
    printf("read header failed.\n");
    fclose(fp);
    unlink(tmp_file);
    return -1;
  }

  printf("alloc hook version %x\n", header.version);

  // Read smaps file
  while (true) {
    std::unique_ptr<SmapsData> smaps_entry(new SmapsData());
    bool ret = FscanfReadFields(fp, "%p-%p %7s %*s %*s %*s", "address_start/address_end/mode", 3,
                                &smaps_entry->address_start, &smaps_entry->address_end, smaps_entry->mode);
    if (!ret) {
      break;
    }

    ret = FscanfReadFields(fp, "%[^\n]\n", "name", 1, smaps_entry->name);   // NOLINT
    if (!ret) {
      break;
    }
    StripSpace(smaps_entry->name);

    ret = FscanfReadFields(fp, "Size: %u %*s\n", "size", 1, &smaps_entry->size);
    if (!ret) {
      break;
    }

    ret = FscanfReadFields(fp, "Rss: %u %*s\n", "rss", 1, &smaps_entry->rss);
    if (!ret) {
      break;
    }
    const char* skip_patterns[] = {"Pss: %*s %*s\n",
                                   "Shared_Clean: %*s %*s\n",
                                   "Shared_Dirty: %*s %*s\n",
                                   "Private_Clean: %*s %*s\n",
                                   "Private_Dirty: %*s %*s\n",
                                   "Referenced: %*s %*s\n",
                                   "Anonymous: %*s %*s\n",
                                   "AnonHugePages: %*s %*s\n",
                                   "Shared_Hugetlb: %*s %*s\n",
                                   "Private_Hugetlb: %*s %*s\n",
                                   "Swap: %*s %*s\n",
                                   "SwapPss: %*s %*s\n",
                                   "KernelPageSize: %*s %*s\n",
                                   "MMUPageSize: %*s %*s\n",
                                   "Locked: %*s %*s\n",
                                   "VmFlags:%*[^\n]\n"};

    for (size_t i = 0; i < sizeof(skip_patterns) / sizeof(skip_patterns[0]); i++) {
      ret = FscanfReadFields(fp, skip_patterns[i], "no_need", 1);
      if (!ret) {
        break;
      }
    }

    if ((strstr(smaps_entry->name, "lib") != nullptr) && (strstr(smaps_entry->name, ".so") != nullptr)) {
      const char* lib_name = basename(smaps_entry->name);
      (void)snprintf(buff, sizeof(buff), "find %s -name %s", lib_dir, lib_name);
      FILE* lib_fp = popen(buff, "r");
      if (lib_fp != nullptr) {
        ret = FscanfReadFields(lib_fp, "%511s", "lib_path", 1, smaps_entry->lib_path);
        if (ret) {
          smaps_entry->is_lib_found = 1;
        }
        (void)pclose(lib_fp);
      }
    }

    if (smaps_list != nullptr) {
      SmapsData* p = smaps_list;
      while (p->next != nullptr) {
        p = p->next;
      }
      p->next = smaps_entry.release();
    } else {
      smaps_list = smaps_entry.release();
    }
  }
  elf_text = smaps_list;

  fseek(fp, sizeof(header) + header.smaps_file_len, SEEK_SET);  // NOLINT

  // Read thread info
  for (int i = 0; i < (header.thread_info_len / sizeof(ThreadRecord)); ++i) {
    auto thread = std::make_unique<ThreadRecord>();
    if (fread(thread.get(), 1, sizeof(ThreadRecord), fp)) {
      thread->next = thread_info;
      thread_info = thread.release();
    } else {
      break;
    }
  }

  // Read alloc list
  SetStackInfo();

  fseek(fp, sizeof(header) + header.smaps_file_len + header.thread_info_len, SEEK_SET);  // NOLINT
  AllocRecord entry;
  while (fread(&entry, 1, sizeof(entry), fp)) {
    AllocRecordParse(&entry);
  }
  fclose(fp);
  unlink(tmp_file);

  PrintHeapResult();   // Print heap result
  PrintMmapResult();   // Print mmap result
  PrintSmapsResult();  // Print smaps result
  CleanList();

  return 0;
}
